import os
import shutil
import tkinter as tk
from tkinter import filedialog



root = tk.Tk()
root.withdraw()


origin_folder_path = filedialog.askdirectory()
print('origin path', origin_folder_path)

merged_folder_path = filedialog.askdirectory()
print('merged path', merged_folder_path)

train_path = os.path.join(merged_folder_path, 'train')
val_path = os.path.join(merged_folder_path, 'validation')
os.mkdir(train_path)
os.mkdir(val_path)

merged_train_xml_path = os.path.join(merged_folder_path, 'train', 'Annotations')
merged_train_img_path = os.path.join(merged_folder_path, 'train', 'JPEGImages')
os.mkdir(merged_train_xml_path)
os.mkdir(merged_train_img_path)


merged_val_xml_path = os.path.join(merged_folder_path, 'validation', 'Annotations')
merged_val_img_path = os.path.join(merged_folder_path, 'validation', 'JPEGImages')
os.mkdir(merged_val_xml_path)
os.mkdir(merged_val_img_path)

print("***************************")

name_index = 0
img_index = 0
xml_index = 0

for root, dirs, files in os.walk(origin_folder_path):
    # root.replace('/', '\\')
    dir_names = root.split('\\')
    dir_name = dir_names[-1]
    print(root)
    print(dir_names)
    print(dir_name)
    print(dirs)
    print(files)

    if dir_name == 'xml':
        if 'train' in dir_names:
            for file_name in files:
                file_path = os.path.join(root, file_name)
                # dst = merged_train_img_path + str(xml_index)
                dst = os.path.join(merged_train_xml_path, str(xml_index) + '.xml')
                shutil.copy(file_path, dst)
                xml_index += 1
            print('copy xml')
        elif 'val' or 'eval' in dir_names:
            for file_name in files:
                file_path = os.path.join(root, file_name)
                # dst = merged_train_img_path + str(xml_index)
                dst = os.path.join(merged_val_xml_path, str(xml_index) + '.xml')
                shutil.copy(file_path, dst)
                xml_index += 1
            print('copy xml')

    elif dir_name == 'img':
        if 'train' in dir_names:
            for file_name in files:
                file_path = os.path.join(root, file_name)
                dst = merged_train_img_path + str(img_index)
                dst = os.path.join(merged_train_img_path, str(img_index) + '.jpg')
                shutil.copy(file_path, dst)
                img_index += 1
            print('copy img')
        elif 'val' or 'eval' in dir_names:
            for file_name in files:
                file_path = os.path.join(root, file_name)
                dst = merged_train_img_path + str(img_index)
                dst = os.path.join(merged_val_img_path, str(img_index) + '.jpg')
                shutil.copy(file_path, dst)
                img_index += 1
            print('copy xml')

    print('===================================')

print('Done')
